import gurobipy as gp
from gurobipy import GRB
import json


def read_data():
    pass



def buildModel(group_num, day_num, time_interval_num, robust_index, demond, demond_t, cost_list, R, U,
            X_B, Y_B, Z_B_1, Z_B_0):
    """ 班组数， 周期长度， 时间间隔个数， 鲁棒性参数， 需求矩阵，d_t需求波动矩阵， 成本列表， 总人数， 缺人代价,
        早班班制， 中班班制，滚动班制 """
    # 上班成本
    C = {} 

    # 早班变量
    x_index = {}
    x_b_index = {}
    # 中班变量
    y_index = {}
    y_b_index = {}

    # 滚动班变量
    z_index = {}
    z_b_index = {}


    # 需求变量
    d_index = {}
    d_t_index = {}
    delta_index = {}
    # 供给变量
    s_index = {}
    # 鲁棒变量
    t_index = {}     #鲁棒系数（j, k）
    
    for t in range(3):
        C[t] = cost_list[t] # 成本系数

    # 生成决策变量
    for i in range(group_num):# i组
        for j in range(day_num):  #j天
            x_index[i,j] = 0  #i班j天早班人数
            y_index[i,j] = 0  #i班j天早班人数
            z_index[i,j,0] = 0  #i班j天滚动白班人数
            z_index[i,j,1] = 0  #i班j天滚动夜班人数

            # 导入班制矩阵
            x_b_index[i,j] = 0
            y_b_index[i,j] = 0
            z_b_index[i,j,1] = 0 # 滚动白班上班情况
            z_b_index[i,j,0] = 0 # 滚动夜班上班情况

    # 生成供给和需求变量s,d,d_t
    for j in range(day_num):
        for k in range(time_interval_num):
            s_index[(j,k)] = 0 # 生成j天k时段的需求
            d_t_index[(j,k)] = demond_t[j][k]    # 需求变动    
            d_index[(j,k)] = demond[j][k]  # 需求均值
            t_index[(j,k)] = robust_index[j][k]  # 鲁棒系数
            delta_index[j,k] = 0 # 每个时段的供需差值

    try:
        m = gp.Model("RO_schedule")
        
        # 设置变量

        x = m.addVars(x_index.keys(),  vtype=GRB.INTEGER, name='x') #早班上班人数（i，j）
        x_b = m.addVars(x_b_index.keys(), vtype=GRB.BINARY, name = "x_b") # 0,1 变量

        y = m.addVars(y_index.keys(), vtype=GRB.INTEGER, name='y' ) #中班上班人数（i，j）
        y_b = m.addVars(y_b_index.keys(), vtype=GRB.BINARY, name = "y_b")  # 0,1 变量

        z = m.addVars(z_index.keys(), vtype=GRB.INTEGER, name='z') #轮班上班人数（i，j，1）（i，j，0）
        z_b = m.addVars(z_b_index.keys(), vtype=GRB.BINARY, name = "z_b")  # 0,1 变量
        
        # 两个可调超参
        # U = m.addVar(lb = U , ub = U) #缺人成本
        # R = m.addVar(lb = R, ub = R) #总职工人数

        delta = m.addVars(delta_index.keys(), vtype=GRB.CONTINUOUS, name = "delta")   #
        lack =m.addVar(vtype=GRB.CONTINUOUS, name = "lack")

        d = m.addVars(d_index.keys(), vtype=GRB.INTEGER, name='d') #需求（j，k）
        d_t = m.addVars(d_t_index.keys(), vtype=GRB.INTEGER, name='d_t') #需求波动（j，k）
        
        t = m.addVars(t_index.keys(), vtype=GRB.CONTINUOUS, name = "t") #鲁棒系数（j, k）
        
        s = m.addVars(s_index.keys(), vtype=GRB.INTEGER, name = "s") #j天k时段的上班总人数（j，k）
        
        # 写入目标函数
        m.setObjective(gp.quicksum(x[i,j]*C[0] +y[i,j]*C[1]+z[i,j,1]*C[2]+z[i,j,0]*C[2] for i in range(group_num) for j in range(day_num)) + U * lack ,GRB.MINIMIZE)

        # 写入约束

        # 人员约束
        m.addConstr((gp.quicksum(x[i,j]+y[i,j]+z[i,j,1]+z[i,j,0] for i in range(group_num) for j in range(day_num))) <= (2*R), "R_constr")
        # 班制约束
        for i in range(group_num):
            for j in range(day_num):
                m.addConstr(x_b[i,j] == X_B[i][j])
                m.addConstr(y_b[i,j] == Y_B[i][j])
                m.addConstr(z_b[i,j,1] == Z_B_1[i][j])
                m.addConstr(z_b[i,j,0] == Z_B_0[i][j])




        for i in range(group_num):
            for j in range(day_num):
                m.addConstr((x_b[i,j] == 1)>>(x[i,j] >= 1))
                m.addConstr((y_b[i,j] == 1)>>(y[i,j] >= 1))
                m.addConstr((z_b[i,j,1] == 1)>>(z[i,j,1] >= 1))
                m.addConstr((z_b[i,j,0] == 1)>>(z[i,j,0] >= 1))

        for i in range(group_num):
            for j in range(day_num):
                m.addConstr((x_b[i,j] == 0)>>(x[i,j] == 0))
                m.addConstr((y_b[i,j] == 0)>>(y[i,j] == 0))
                m.addConstr((z_b[i,j,1] == 0)>>(z[i,j,1] == 0))
                m.addConstr((z_b[i,j,0] == 0)>>(z[i,j,0] == 0))
        
        # 供给约束
        for j in range(day_num):
            m.addConstr(s[j,0] == gp.quicksum(x[i,j] for i in range(group_num))) # 第j天的第一时段的供给
            m.addConstr(s[j,1] == (gp.quicksum(x[i,j] for i in range(group_num)) + (gp.quicksum(z[i,j,1] for i in range(group_num)))))
            m.addConstr(s[j,2] == (gp.quicksum(y[i,j] for i in range(group_num)) + (gp.quicksum(z[i,j,1] for i in range(group_num)))))
            m.addConstr(s[j,3] == (gp.quicksum(y[i,j] for i in range(group_num)) + (gp.quicksum(z[i,j,0] for i in range(group_num)))))
            m.addConstr(s[j,4] == gp.quicksum(z[i,j,0] for i in range(group_num)))
        

        # 需求约束 给需求赋值
        for j in range(day_num):
            for k in range(day_num):
                m.addConstr(d[j,k] == ((demond[j][k] + demond_t[j][k])-s[j,k]))

        # 供需差值约束
        for j in range(day_num):
            for k in range(time_interval_num):
                # m.addGenConstrMax(delta[j,k], [s[j,k] - d[j,k], 0])

                m.addConstr(delta[j,k] == (gp.max_([d[j,k], 0])), "delta_constr")
                aaa = 1
        m.addConstr(lack == gp.quicksum(delta[j,k] for j in range(day_num) for k in range(time_interval_num)))

        # 求解
        m.optimize()

        #输出决策变量
        for v in m.getVars():
            print("{}:{}".format(v.varName, v.x))
            # print(v.varName, v.x)
        # 输出目标
        print("目标函数值：{}".format(m.objVal))

        total_x = 0
        total_y = 0
        total_z = 0

        for i in x_index.keys():
            total_x += x[i].x
        for i in y_index.keys():
            total_y += y[i].x
        for i in z_index.keys():
            total_z += z[i].x

        print([s[k].x for k in s_index.keys()])


        print(total_x, total_y, total_z)
        


    except gp.GurobiError as e:
        print("Error code  "+ str(e.errno) + ": "+ str(e))
    except AttributeError:
        print("encontered an attribute error")
    
if __name__ == "__main__":
    # 读数据
    data = read_data()
    
    # 写入参数值
    group_num = 4
    day_num =  4
    time_interval_num = 5
    robust_index = [[1,1,1,1,1],
                    [1,1,1,1,1],
                    [1,1,1,1,1],
                    [1,1,1,1,1]]
    cost_list = [500,500,500]

    demond = [[43,179,214,207,193],
              [73,184,206,216,201],
              [41,153,174,177,161],
              [74,179,217,225,215]]

    demond_t = [[0,0,0,0,0],
                [0,0,0,0,0],
                [0,0,0,0,0],
                [0,0,0,0,0]]

    R = 7000 # 员工数
    U = 5000# 缺人成本

    # 班制
    X_B = [[1,1,0,0],[0,1,1,0],[0,0,1,1],[1,0,0,1]]
    Y_B = [[1,1,0,0],[0,1,1,0],[0,0,1,1],[1,0,0,1]]
    Z_B_1 = [[1,0,0,0],[0,1,0,0],[0,0,1,0],[0,0,0,1]]
    Z_B_0 = [[0,1,0,0],[0,0,1,0],[0,0,0,1],[1,0,0,0]]

    buildModel(group_num, day_num, time_interval_num, robust_index, demond, demond_t, cost_list, R, U, X_B, Y_B, Z_B_1, Z_B_0)
